import os
import sys
import argparse
import polib
import asyncio
from tqdm import tqdm

# 将项目根目录添加到 Python 路径，以便导入 src 中的模块
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, project_root)

from src.translation.openai_translator import OpenAITranslator, TranslationError
from src.translation.microsoft_translator import MicrosoftTranslator
from src.config import app_config


# 语言代码映射 (locale dir -> language name for API)
# 关键点: 字典的键 (key) 必须与 'locales/' 目录下的文件夹名完全一致。
# 值 (value) 是翻译 API 能理解的语言名称或代码。
LANGUAGE_MAP = {
    'openai': {
        'ja-JP': 'Japanese',
        'fr-FR': 'French',
        'de-DE': 'German',
        'ru-RU': 'Russian',
        'zh-CN': 'Chinese (Simplified)',
        'zh-TW': 'Chinese (Traditional)',
    },
    'microsoft': {
        'ja-JP': 'ja',
        'fr-FR': 'fr',
        'de-DE': 'de',
        'ru-RU': 'ru',
        'zh-CN': 'zh-Hans',
        'zh-TW': 'zh-Hant',
    }
}

SOURCE_LANGUAGE = "en"


def get_translator(service_name):
    """根据服务名称获取翻译器实例和配置"""
    if service_name == 'microsoft':
        print("Initializing Microsoft Translator...")
        return MicrosoftTranslator(), None
    
    elif service_name == 'openai':
        print("Initializing OpenAI Translator...")
        try:
            translator = OpenAITranslator()
            service_providers = app_config.get('service_providers', [])
            if not service_providers:
                raise ValueError("No 'service_providers' configured in config.yaml.")

            for provider in service_providers:
                if not (provider.get('services') and provider['services'].get('llm') and provider['services']['llm'].get('models')):
                    continue

                models = provider['services']['llm']['models']
                for model in models:
                    if 'GLM' in model.get('model_name', ''):
                        model_config = {
                            "api_key": provider.get("api_key"),
                            "base_url": provider['services']['llm'].get("base_url"),
                            "model_name": model.get("model_name"),
                            "display_name": model.get("display_name"),
                        }
                        translator._get_client_for_model(model_config)
                        print(f"Found compatible model: {model_config['display_name']}")
                        return translator, model_config
            
            raise ValueError("No model containing 'GLM' found in any service provider.")

        except (KeyError, IndexError, ValueError) as e:
            print(f"Error: Could not initialize translator or get model config. Please check your config.yaml.")
            print(f"Details: {e}")
            sys.exit(1)
    else:
        print(f"Error: Unknown service '{service_name}'.")
        sys.exit(1)


async def main():
    """主函数，执行翻译流程"""
    parser = argparse.ArgumentParser(description="使用 AI 翻译 .po 文件。")
    parser.add_argument(
        '-l', '--languages',
        nargs='+',
        help="要翻译的语言代码列表 (例如 'fr-FR' 'ja-JP')。如果未提供，则处理所有语言。"
    )
    parser.add_argument(
        '--service',
        type=str,
        default='microsoft',
        choices=['microsoft', 'openai'],
        help="要使用的翻译服务 (默认: microsoft)。"
    )
    args = parser.parse_args()

    locales_dir = os.path.join(project_root, 'locales')

    print(f"--- Using translation service: {args.service} ---")
    translator, model_config = get_translator(args.service)

    # 确定要处理的语言列表
    if args.languages:
        locales_to_process = args.languages
        print(f"\n--- 将只处理指定的语言: {', '.join(locales_to_process)} ---")
    else:
        locales_to_process = os.listdir(locales_dir)
        print("\n--- 将处理 locales 目录下的所有语言 ---")

    # --- 循环处理所有语言 ---
    for locale_name in locales_to_process:
        if not os.path.isdir(os.path.join(locales_dir, locale_name)) or locale_name == 'en-US':
            continue

        service_lang_map = LANGUAGE_MAP.get(args.service, {})
        target_language = service_lang_map.get(locale_name)
        
        if not target_language:
            print(f"Warning: No language mapping found for '{locale_name}' with service '{args.service}'. Skipping.")
            continue

        po_file_path = os.path.join(
            locales_dir, locale_name, 'LC_MESSAGES', 'messages.po')
        if not os.path.exists(po_file_path):
            print(f"警告：在 '{locale_name}' 目录中未找到 messages.po 文件，已跳过。")
            continue

        print(f"\n--- 正在处理语言: {target_language} ({locale_name}) ---")
        po = polib.pofile(po_file_path)

        # 强制重新翻译所有条目
        entries_to_translate = [entry for entry in po if entry.msgid]

        if not entries_to_translate:
            print("没有找到需要翻译的条目。")
            continue

        print(f"找到 {len(entries_to_translate)} 个条目进行翻译。")

        original_texts = [entry.msgid for entry in entries_to_translate]

        try:
            # 使用 tqdm 创建进度条
            with tqdm(total=len(original_texts), desc=f"Translating to {locale_name}") as pbar:

                def progress_callback(done, total, batch_size):
                    pbar.update(batch_size)
                
                # 根据翻译器类型决定调用方式
                if isinstance(translator, OpenAITranslator):
                    translated_texts = await asyncio.to_thread(
                        translator.translate,
                        texts=original_texts,
                        source_lang=SOURCE_LANGUAGE,
                        target_lang=target_language,
                        model_config=model_config,
                        batch_done_callback=progress_callback
                    )
                elif isinstance(translator, MicrosoftTranslator):
                    # MicrosoftTranslator is not async, but we run it in a thread
                    # It handles batching internally, so no callback is needed.
                    translated_texts = await asyncio.to_thread(
                        translator.translate,
                        texts=original_texts,
                        source_lang=SOURCE_LANGUAGE,
                        target_lang=target_language
                    )
                    pbar.update(len(original_texts)) # Manually update progress bar
                else:
                    raise TypeError(f"Unsupported translator type: {type(translator)}")


            if len(translated_texts) != len(original_texts):
                raise TranslationError(
                    f"Translation returned {len(translated_texts)} items, but expected {len(original_texts)}.")

            for entry, translation in zip(entries_to_translate, translated_texts):
                entry.msgstr = translation

            print(f"正在保存文件: {po_file_path}")
            po.save()
            print("保存成功。")

        except TranslationError as e:
            print(f"\n错误：翻译 '{target_language}' 时发生错误。")
            print(e)
        except Exception as e:
            print(f"\n错误：处理 '{locale_name}' 时发生未知错误: {e}")

if __name__ == "__main__":
    # 检查配置文件是否存在
    if not os.path.exists(os.path.join(project_root, 'config.yaml')):
        print("错误: config.yaml 文件未找到。请先从 config.yaml.example 创建并配置它。")
        sys.exit(1)

    asyncio.run(main())
